'''
'''
#!/usr/bin/env python 
# -*- encoding: utf-8 -*-
'''
Desc: 
Author: KouXichao
Date: 2020-12-21 21:37:08
Contact: kouxichao@qq.com
LastEditors: KouXichao
LastEditTime: 2021-01-18 16:00:28
FilePath: /LayoutGen/gendata.py
Github: https://github.com/kouxichao
-----------------------------------------------------------------------------
FreeType high-level python API - Copyright 2020-2020 Nicolas P. Rougier
Distributed under the terms of the new BSD license.
-----------------------------------------------------------------------------

'''


import argparse
import random
import shutil
import os
from randomlayout import randomlayout, ConfigGen
from main import main
from create_xml import make_xml
from data2labelme_json import make_json
from labelme2coco import labelme2coco
import glob
import copy

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = "LayoutGen !!!")
    # image options
    parser.add_argument(
        '-lan',
        "--language", 
        default='en', 
        type=str, 
        choices=['en', 'jp', 'cn'], 
        help='Which language to render!')
    parser.add_argument(
        '-rl',
        "--random_lan", 
        default=False, 
        action="store_true",
        help='Random choose language to render! this option will surpress -lan opt')
    parser.add_argument(
        "--width",
        default=960,
        type=int,
        help="Width of image"
    )
    parser.add_argument(
        "--height",
        default=1280,
        type=int,
        help="Height of image"
    )
    parser.add_argument(
        '-rb',
        "--random_bg",
        default=False,
        action="store_true",
        help="Whether to randomly select the background image, the default background is white"
    )
    parser.add_argument(
        '--scale', 
        default=1.0, 
        type=float, 
        help="scale value of resizing generated image and boxes!")
    parser.add_argument(
        '--warp', 
        default=False, 
        action='store_true', 
        help="warp image.")
    parser.add_argument(
        '-ar',
        "--angle_range", 
        default=30, 
        type=int, 
        help="Range of Warp angle")
    parser.add_argument(
        "-tc",
        "--thresh_color", 
        default=(0, 80), 
        type=tuple, 
        help="Define color diff val between bg and fg")
        
    # font options
    parser.add_argument(
        '-ft', 
        "--figure_type", 
        default=0, 
        type=int, 
        help="Types of figures to be inserted!")
    parser.add_argument(
        "-fs",
        "--font_size", 
        default=(20, 30, None), 
        type=tuple, 
        help="Define font size of text")

    # Layout options
    parser.add_argument(
        '-th',
        "--text_height",
        default=0.02,
        type=float,
        help="min height of text region"
    )
    parser.add_argument(
        '-tr',
        "--type_ratio",
        default=[4, 1, 1, 0, 2], # text:image:graph:table:tableimage
        type=list,
        help="Area type ratio of pages"
    )
    parser.add_argument(
        '-nc',
        "--num_column",
        default=3,
        type=int,
        help="Random range of the number of vertical columns on the page"
    )
    parser.add_argument(
        '-nt',
        "--num_title",
        default=3,
        type=int,
        help="Random range of the number of title on the page"
    )
    parser.add_argument(
        '-nf',
        "--num_formula",
        default=5,
        type=int,
        help="Random range of the number of formula on the page"
    )

    # Data options
    parser.add_argument(
        '-dd',
        '--dateSaveDir',
        default="./synthetic_data/",
        type=str,
        help="Path of saving generated images and gts"
    )    
    parser.add_argument(
        '-nl', 
        "--num_layout", 
        default=1,
        type=int, 
        required=True, 
        help="The number of layout to generate!")
    parser.add_argument(
        '-ni', 
        "--num_img", 
        default=1,
        type=int, 
        help="The number of images generated by the same layout!")
    parser.add_argument(
        '-rm', 
        "--remove_data", 
        default=False, 
        action="store_true", 
        help="Whether to remove the generated data!")

    # Debug and develpment options
    parser.add_argument(
        '-sl',
        "--save_layout",
        default=False,
        action="store_true",
        help="Whether to save layout information in the form of pictures, the default is saved in the profile_image folder"
    )
    parser.add_argument(
        '-vis',
        "--debug",
        default=False,
        action="store_true",
        help="Whether to save layout information in the form of pictures, the default is saved in the profile_image folder"
    )
    
    opt = parser.parse_args()
    
    # Remove previous generated data!!!
    if opt.remove_data and os.path.exists(opt.dateSaveDir):
        print("Warning !!!\nremove all data in {}(y/n)?".format(opt.dateSaveDir))
        inp = input()
        if inp == 'y':
            shutil.rmtree(opt.dateSaveDir)

        # os.mkdir(opt.dateSaveDir)
    xml_path = os.path.join(opt.dateSaveDir, 'xml/')
    json_path = os.path.join(opt.dateSaveDir, 'json/')
    language = opt.language
    # Generate pages
    for i in range(opt.num_layout):
        # Layout generation
        if opt.random_lan:
            language = random.choice(['en', 'cn', 'jp'])
        vcolumn = random.randint(1,opt.num_column)
        ntitle = random.randint(0,opt.num_title)
        nformula = 0
        if vcolumn < 3:
            nformula = random.randint(0, opt.num_formula)
        rectlist = randomlayout(opt.width, 
                            opt.height, 
                            language, 
                            vcolumn, 
                            ntitle, 
                            nformula, 
                            opt.text_height,
                            debug = opt.debug,
                            type_ratio = opt.type_ratio) # text:image:graph:table:tableimage
        page_config_ = ConfigGen(opt.width, 
                                opt.height,
                                language,
                                rectlist[0], 
                                rectlist[1], 
                                rectlist[2], 
                                rectlist[3], 
                                rectlist[4], 
                                rectlist[5], 
                                rectlist[6],
                                rectlist[7])  
        for j in range(opt.num_img):
            page_config = copy.deepcopy(page_config_)
            nth = opt.num_img * i + j + 1
            print("\033[31m**********************************************")
            print("        start to generate %dth image          "%(nth))
            print("**********************************************\033[0m")
            # Render image
            opt.font_size = (opt.font_size[0], opt.font_size[1], random.randint(opt.font_size[0], opt.font_size[1])) 
            page_config = main(opt.font_size, 
                opt.scale, 
                opt.figure_type, 
                opt.thresh_color,
                opt.warp, 
                opt.angle_range, 
                opt.dateSaveDir,
                page_config=page_config,
                random_bg=opt.random_bg,
                save_layout=opt.save_layout)  

            # GTs of VOC(xml) format
            if not os.path.exists(xml_path):
                os.mkdir(xml_path)
            make_xml(xml_path, page_config)

            # GTs of COCO(json) format
            if not os.path.exists(json_path):
                os.mkdir(json_path)
            make_json(json_path, page_config)
    
    # GT of COCO(json) format
    labelme_json = glob.glob(json_path+'[0-9]*_[0-9]*.json')
    outpath = os.path.join(json_path, 'instances_train2014.json')
    labelme2coco(labelme_json, outpath)


